/*
 * Copyright 2010 Google Inc.
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package com.google.websocket;

import com.google.util.Base64Utils;


import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.AlgorithmParameterSpec;
import java.util.Arrays;
import java.util.Map;
import java.util.Random;
import java.util.Set;

import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.CipherOutputStream;
import javax.crypto.Mac;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

/**
 * Implementation of Adam Barth's proposed WebSocket handshake.
 * 
 * The handshake is described at
 * http://www.ietf.org/id/draft-abarth-websocket-handshake-00.txt
 */
public class AdamBarthHandshake extends HttpCompatibleHandshake {

  private static final String BOGUS_HOST
      = "1C1BCE63-1DF8-455C-8235-08C2646A4F21.invalid:443";
  private static final String HTTP_RESPONSE = "HTTP/1.1 200 OK\r\n";
  private static final String WS_CONNECT = "CONNECT " + BOGUS_HOST
      + " HTTP/1.1\r\n";
  private static final String WS_HOST = "Host: " + BOGUS_HOST + "\r\n";
  private static final String WS_KEY = "Sec-WebSocket-Key: ";

  private static final String CLIENT_HANDSHAKE_KEY
      = "C1BA787A-0556-49F3-B6AE-32E5376F992B";
  private static final String SERVER_HANDSHAKE_KEY
      = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
  private static final String CLIENT_TO_SERVER_KEY
      = "363A6078-74D2-4C0B-8CBC-1E6A36E83442";
  private static final String SERVER_TO_CLIENT_KEY
      = "2306C3BE-0ACF-42C0-B69E-DFFE02CFA346";
  
  // performs steps 6-27 of section 1.1 in Adam Barth's spec
  public WebSocket clientHandshake(InputStream in, OutputStream out,
      Set<Extension> extensions, URI uri, String origin) throws IOException {
    Random rng = new SecureRandom();
    byte[] clientNonce = new byte[16];
    rng.nextBytes(clientNonce);
    out.write(WS_CONNECT.getBytes(UTF8));
    out.write(WS_HOST.getBytes(UTF8));
    String keyHdr = WS_KEY + Base64Utils.toBase64(clientNonce) + "\r\n";
    out.write(keyHdr.getBytes(UTF8));
    out.write("\r\n".getBytes(UTF8));
    byte[] handshakeKey = getKey(CLIENT_HANDSHAKE_KEY, clientNonce);
    OutputStream maskedOut = maskOutputStream(out, handshakeKey);
    try {
      uri = normalizeUri(uri);
    } catch (URISyntaxException e) {
      throw new IOException("invalid URI", e);
    }
    String path = uri.getPath();
    if (!path.startsWith("/")) {
      path = "/" + path;
    }
    String getLine = "GET " + path + " HTTP/1.1\r\n";
    maskedOut.write(getLine.getBytes(UTF8));
    String host = uri.getHost();
    int port = uri.getPort();
    String hostLine = "Host: " + host + ":" + port + "\r\n";
    maskedOut.write(hostLine.getBytes(UTF8));
    String originLine = "Origin: " + origin + "\r\n";
    maskedOut.write(originLine.getBytes(UTF8));
    // TODO(jat): cookies
    maskedOut.write("\r\n".getBytes(UTF8));
    maskedOut.flush();
    String line = Util.readUtf8Line(in);
    if (!HTTP_RESPONSE.equals(line + "\r\n")) {
      throw new IOException("Unexpected HTTP response");
    }
    Map<String, String> headerFields = Util.readHeaderFields(in);
    String serverAccept = headerFields.get("sec-websocket-accept");
    String serverOrigin = headerFields.get("sec-websocket-origin");
    String serverKey = headerFields.get("sec-websocket-key");
    String serverLocation = headerFields.get("sec-websocket-location");
    byte[] serverAcceptHash = Base64Utils.fromBase64(serverAccept);
    byte[] handshakeResponseKey = getKey(SERVER_HANDSHAKE_KEY, clientNonce);
    if (!Arrays.equals(serverAcceptHash, handshakeResponseKey)) {
      throw new IOException("response key mismatch");
    }
    byte[] serverNonce = Base64Utils.fromBase64(serverKey);
    if (serverNonce.length != 16 || Arrays.equals(clientNonce, serverNonce)) {
      throw new IOException("bad server nonce");
    }
    if (!origin.equals(serverOrigin)) {
      throw new IOException("origin mismatch: server=" + serverOrigin
          + ", mine=" + origin);
    }
    if (!serverLocation.equals(uri.toASCIIString())) {
      throw new IOException("location mismatch: server=" + serverLocation
          + ", client=" + uri.toASCIIString());
    }
    // TODO(jat): cookies
    byte[] combinedNonce = new byte[32];
    System.arraycopy(clientNonce, 0, combinedNonce, 0, 16);
    System.arraycopy(serverNonce, 0, combinedNonce, 16, 16);
    byte[] csKey = getKey(CLIENT_TO_SERVER_KEY, combinedNonce);
    byte[] scKey = getKey(SERVER_TO_CLIENT_KEY, combinedNonce);
    return new DraftWebSocket(maskInputStream(in, scKey), maskOutputStream(out,
        csKey), extensions);
  }


  @Override
  public boolean matches(String line, Map<String, String> headers) {
    if (!WS_CONNECT.equals(line + "\r\n")) {
      return false;
    }
    if (!BOGUS_HOST.equals(headers.get("host"))) {
      return false;
    }
    if (!headers.containsKey("sec-websocket-key")) {
      return false;
    }
    return true;
  }

  @Override
  public WebSocket serverHandshake(String line,
      Map<String, String> topHeaders, InputStream in, OutputStream out,
      Set<Extension> extensions, int port) throws IOException {
    String key = topHeaders.get("sec-websocket-key");
    byte[] clientNonce = Base64Utils.fromBase64(key);
    byte[] clientHandshakeKey = getKey(CLIENT_HANDSHAKE_KEY, clientNonce);
    InputStream maskedIn = unmaskInputStream(in, clientHandshakeKey);
    line = Util.readUtf8Line(maskedIn);
    Map<String, String> headerFields = Util.readHeaderFields(maskedIn);
    String[] get = line.split(" ");
    if (get.length != 3 || !"GET".equals(get[0])) {
      throw new IOException("Incorrect WebSocket handshake");
    }
    String resource = get[1];
    // TODO(jat); validate resource
    String host = headerFields.get("host");
    // TODO(jat): validate host
    String origin = headerFields.get("origin");
    // TODO(jat): validate origin
    String proto = headerFields.get("sec-websocket-protocol");
    if (proto != null) {
      throw new IOException("Unsupported protocol '" + proto + "' requested");
    }
    Random rng = new SecureRandom();
    byte[] serverNonce = new byte[16];
    rng.nextBytes(serverNonce);
    byte[] handshakeResponseKey = getKey(SERVER_HANDSHAKE_KEY, clientNonce);
    out.write(HTTP_RESPONSE.getBytes(UTF8));
    line = "Sec-WebSocket-Accept: "
        + Base64Utils.toBase64(handshakeResponseKey) + "\r\n";
    out.write(line.getBytes(UTF8));
    line = "Sec-WebSocket-Key: "
        + Base64Utils.toBase64(serverNonce) + "\r\n";
    out.write(line.getBytes(UTF8));
    line = "Sec-WebSocket-Origin: " + origin + "\r\n";
    out.write(line.getBytes(UTF8));
    URI uri;
    try {
      uri = new URI("ws://" + host + resource);
    } catch (URISyntaxException e) {
      throw new IOException("invalid URI", e);
    }
    line = "Sec-WebSocket-Location: " + uri.toASCIIString() + "\r\n";
    out.write(line.getBytes(UTF8));
    out.write("\r\n".getBytes(UTF8));
    byte[] combinedNonce = new byte[32];
    System.arraycopy(clientNonce, 0, combinedNonce, 0, 16);
    System.arraycopy(serverNonce, 0, combinedNonce, 16, 16);
    byte[] csKey = getKey(CLIENT_TO_SERVER_KEY, combinedNonce);
    byte[] scKey = getKey(SERVER_TO_CLIENT_KEY, combinedNonce);
    return new DraftWebSocket(maskInputStream(in, csKey), maskOutputStream(out,
        scKey), extensions);
  }

  protected Cipher getCipher(byte[] key, boolean encrypt) throws IOException {
    Throwable caught = null;
    try {
      Cipher cipher = Cipher.getInstance("AES/CTR/NoPadding");
      // Note that using the key for the IV is not secure, but in this case we
      // are transmitting the key in the clear anyway since we only want to make
      // sure an attacker can't control bytes in the handshake.
      AlgorithmParameterSpec paramSpec = new IvParameterSpec(key);
      SecretKeySpec secretKeySpec = new SecretKeySpec(key, "AES");
      cipher.init(encrypt ? Cipher.ENCRYPT_MODE : Cipher.DECRYPT_MODE,
          secretKeySpec, paramSpec);
      return cipher;
    } catch (InvalidKeyException e) {
      caught = e;
    } catch (InvalidAlgorithmParameterException e) {
      caught = e;
    } catch (NoSuchAlgorithmException e) {
      caught = e;
    } catch (NoSuchPaddingException e) {
      caught = e;
    }
    throw new IOException("Unable to setup encryption", caught);
  }

  protected InputStream maskInputStream(InputStream istr, byte[] key)
      throws IOException {
    return new CipherInputStream(istr, getCipher(key, true));
  }

  protected OutputStream maskOutputStream(OutputStream ostr, byte[] key)
      throws IOException {
    return new CipherOutputStream(ostr, getCipher(key, true));
  }

  protected InputStream unmaskInputStream(InputStream istr, byte[] key)
      throws IOException {
    return new CipherInputStream(istr, getCipher(key, true));
  }

  protected OutputStream unmaskOutputStream(OutputStream ostr, byte[] key)
      throws IOException {
    return new CipherOutputStream(ostr, getCipher(key, true));
  }

  /**
   * Generates a AES key by taking the first 16 bytes of the HMAC-SHA1 hash of
   * the string, keyed with the supplied nonce.
   * 
   * @param uuid
   * @param nonce
   * @return 16-byte key
   * @throws IOException
   */
  private byte[] getKey(String uuid, byte[] nonce) throws IOException {
    Throwable caught;
    try {
      SecretKey key = new SecretKeySpec(nonce, "HmacSHA1");
      Mac mac = Mac.getInstance("HmacSHA1");
      mac.init(key);
      mac.update(uuid.getBytes(UTF8));
      byte[] hash = mac.doFinal();
      byte[] result = new byte[16];
      System.arraycopy(hash, 0, result, 0, 16);
      return result;
    } catch (InvalidKeyException e) {
      caught = e;
    } catch (NoSuchAlgorithmException e) {
      caught = e;
    } catch (IllegalStateException e) {
      caught = e;
    }
    throw new IOException("Unable to initialize SHA1 HMAC", caught);
  }
}
